--------------------------[ fun_writefasta(text) ]-----------------------------
--
-- Returns a fasta file given a two-column table.
--
-- Input: 
-- a SELECT query returning the sequence name as first column and the sequence 
-- as second column.
-- Adjust variable conn for appropriate connection parameters
--
-- Output: 
-- File '/tmp/writefasta.fa'. See variable outfile below to change this.
--
-- Usage:
-- SELECT fun_writefasta('SELECT fasta_name, fasta_seq FROM myTable');
-- 
-- Note: 
-- If the SELECT query contains single quotes, replace with two single quotes.
-- E.g. SELECT fun_writefasta('SELECT fasta_name, ''actg'' FROM myTable');

-------------------------------------------------------------------------------

-- drop function fun_writefasta(text);

create or replace function fun_writefasta(fasta_qry text) returns text AS $$

# --------------------------[ User's input ]-----------------------------------

outfile= 'C:/Documents and Settings/postgres/My Documents/Tritume/writefasta.fa' ## '/tmp/writefasta.fa'

#------------------------------------------------------------------------------

  import psycopg2
  import os
  import re
  
  try:
      conn = psycopg2.connect("dbname='Soay' user='dberaldi' host='localhost' password='mypassword'");
  except:
      return("I am unable to connect to the database")

  cur = conn.cursor()
  cur.execute(fasta_qry)

  fout= open(outfile, 'w')
  
  row = cur.fetchone()
  if len(row) != 2:
      return('Error: fasta_qry must return exactly two columns.\n' + str(len(row)) + ' column(s) found.') 
      fout.close()
      stop  

  aa= re.compile('[^ACDEFGHIKLMNPQRSTVWYBUXZacdefghiklmnpqrstvwybuxz\n\r]')
  warn= ''
  
  while row != None:
      fasta_name= '>' + str(row[0])
      fasta_seq= row[1]

      if aa.search(row[1]) != None:
          warn= ('Warning: Non fasta character found in sequence...\n')
          
      fout.write(fasta_name + '\n' + fasta_seq + '\n')
      row= cur.fetchone()

  fout.close()

  os.system('chmod 777 '+ outfile)  

  return(warn + 'Fasta file written to ' + outfile)
    
  $$ LANGUAGE plpythonu;

